import psycopg2

from psycopg2 import OperationalError, InterfaceError, connect as pg_connect
from sqlalchemy.engine import create_engine, Engine
from ..utils import HawthornConfig

class DatabaseConnectionManager(HawthornConfig):
    provider:str="postgres"
    def __init__(self, config_file: str):
        super().__init__(config_file)
        self.connection=self.__connect()
        self.engine=self.__create_engine()
    def set_provider(self, provider: str):
        self.provider = provider 
    def connect(self) -> psycopg2.extensions.connection:
        return self.connection
    def __create_engine(self) -> Engine:
        if self.provider == "opengauss":
            from sqlalchemy.dialects.postgresql.base import PGDialect
            PGDialect._get_server_version_info = lambda *args: (9, 2)
            return create_engine(self.dsn, pool_pre_ping=True)
        return create_engine(self.dsn, pool_pre_ping=True)
    def __connect(self) -> psycopg2.extensions.connection:
        """
        连接到PostgreSQL数据库
        
        返回:
            psycopg2连接对象
            
        抛出:
            ConnectionError: 当连接失败时，包含详细错误信息
        """
        try:
            self._connection = pg_connect(**self.config)
            return self._connection
        except OperationalError as e:
            raise ConnectionError(f"连接操作失败: {e}")
        except InterfaceError as e:
            raise ConnectionError(f"连接接口错误: {e}")
        except Exception as e:
            raise ConnectionError(f"连接时发生意外错误: {e}")
    def close(self):
        """关闭数据库连接"""
        if self._connection and not self._connection.closed:
            self._connection.close()
        self._connection = None